#!/usr/bin/env python3
"""
统一图算子精度和性能测试脚本

该脚本用于测试三个核心图算子：
1. edge_softmax - 边注意力权重计算
2. gspmm - 广义稀疏矩阵乘法
3. gsddmm - 广义稀疏-密集-密集矩阵乘法

测试包括：
- 精度对比：PyG vs DGL CPU
- 性能对比：PyG vs DGL CPU
- 大规模性能测试
"""

import os
import sys
import subprocess
import argparse
import re
from pathlib import Path
from typing import Dict, List, Tuple, Optional

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

class TestDataCollector:
    """测试数据收集器"""
    
    def __init__(self):
        self.precision_data = {}
        self.performance_data = {}
        self.large_scale_data = {}
        self.test_results = {}
    
    def parse_test_output(self, test_name: str, output: str):
        """解析测试输出，提取数据"""
        lines = output.split('\n')
        
        # 解析精度数据
        for line in lines:
            if '[edge_softmax 精度]' in line:
                self._parse_precision_line('edge_softmax', line)
            elif '[gspmm copy_e+sum 精度]' in line:
                self._parse_precision_line('gspmm_copy_e', line)
            elif '[gsddmm dot 精度]' in line:
                self._parse_precision_line('gsddmm_dot', line)
            elif '[gsddmm add 精度]' in line:
                self._parse_precision_line('gsddmm_add', line)
        
        # 解析性能数据
        for line in lines:
            if 'edge_softmax 性能对比' in line:
                self._parse_performance_section(lines, 'edge_softmax')
            elif 'gspmm 性能对比' in line:
                self._parse_performance_section(lines, 'gspmm')
            elif 'gsddmm 性能对比' in line:
                self._parse_performance_section(lines, 'gsddmm')
        
        # 解析大规模性能数据
        for line in lines:
            if '大规模 edge_softmax 性能' in line:
                self._parse_large_scale_section(lines, 'edge_softmax')
            elif '大规模 gspmm 性能' in line:
                self._parse_large_scale_section(lines, 'gspmm')
            elif '大规模 gsddmm 性能' in line:
                self._parse_large_scale_section(lines, 'gsddmm')
    
    def _parse_precision_line(self, operator: str, line: str):
        """解析精度数据行"""
        # 提取MAE, MAX_ABS, MAX_REL
        mae_match = re.search(r'MAE=([\d.e+-]+)', line)
        max_abs_match = re.search(r'MAX_ABS=([\d.e+-]+)', line)
        max_rel_match = re.search(r'MAX_REL=([\d.e+-]+)', line)
        
        if mae_match and max_abs_match and max_rel_match:
            self.precision_data[operator] = {
                'mae': float(mae_match.group(1)),
                'max_abs': float(max_abs_match.group(1)),
                'max_rel': float(max_rel_match.group(1))
            }
    
    def _parse_performance_section(self, lines: List[str], operator: str):
        """解析性能数据部分"""
        for i, line in enumerate(lines):
            if f'{operator} 性能对比' in line:
                # 查找PyG和DGL的时间数据
                for j in range(i, min(i+10, len(lines))):
                    if 'PyG:' in lines[j] and 'DGL:' in lines[j+1]:
                        pyg_line = lines[j]
                        dgl_line = lines[j+1]
                        
                        # 提取时间数据
                        pyg_time = self._extract_time(pyg_line)
                        dgl_time = self._extract_time(dgl_line)
                        
                        if pyg_time and dgl_time:
                            self.performance_data[operator] = {
                                'pyg_time': pyg_time,
                                'dgl_time': dgl_time
                            }
                        break
                break
    
    def _parse_large_scale_section(self, lines: List[str], operator: str):
        """解析大规模性能数据部分"""
        for i, line in enumerate(lines):
            if f'大规模 {operator} 性能' in line:
                # 查找总时间和平均时间
                for j in range(i, min(i+10, len(lines))):
                    if '总时间:' in lines[j]:
                        total_time = self._extract_time(lines[j])
                        avg_time = self._extract_time(lines[j+1]) if j+1 < len(lines) else None
                        
                        if total_time:
                            self.large_scale_data[operator] = {
                                'total_time': total_time,
                                'avg_time': avg_time or total_time / 10
                            }
                        break
                break
    
    def _extract_time(self, line: str) -> Optional[float]:
        """从行中提取时间数据"""
        # 匹配各种时间格式
        patterns = [
            r'(\d+\.\d+)s',
            r'(\d+\.\d+)\s*ms',
            r'(\d+\.\d+)\s*ms/iter'
        ]
        
        for pattern in patterns:
            match = re.search(pattern, line)
            if match:
                return float(match.group(1))
        return None
    
    def print_summary_tables(self):
        """打印精简汇总（仅真实数据）"""
        # 精度
        self._print_precision_table()
        # 性能
        self._print_performance_table()
    
    def _print_precision_table(self):
        """打印精度（仅数据）"""
        print("\n精度 (PyG vs DGL CPU):")
        # 只打印存在的数据
        if 'edge_softmax' in self.precision_data:
            d = self.precision_data['edge_softmax']
            print(f"edge_softmax mae={d.get('mae', 0.0):.6e} max_abs={d.get('max_abs', 0.0):.6e} max_rel={d.get('max_rel', 0.0):.6e}")
        if 'gspmm_copy_e' in self.precision_data:
            d = self.precision_data['gspmm_copy_e']
            print(f"gspmm(copy_e) mae={d.get('mae', 0.0):.6e} max_abs={d.get('max_abs', 0.0):.6e} max_rel={d.get('max_rel', 0.0):.6e}")
        if 'gsddmm_dot' in self.precision_data:
            d = self.precision_data['gsddmm_dot']
            print(f"gsddmm(dot) mae={d.get('mae', 0.0):.6e} max_abs={d.get('max_abs', 0.0):.6e} max_rel={d.get('max_rel', 0.0):.6e}")
        if 'gsddmm_add' in self.precision_data:
            d = self.precision_data['gsddmm_add']
            print(f"gsddmm(add) mae={d.get('mae', 0.0):.6e} max_abs={d.get('max_abs', 0.0):.6e} max_rel={d.get('max_rel', 0.0):.6e}")
    
    def _print_performance_table(self):
        """打印性能（仅数据）"""
        print("\n性能 (PYG NPU vs DGL CPU):")
        for operator in ['edge_softmax', 'gspmm', 'gsddmm']:
            data = self.performance_data.get(operator, {})
            pyg_time = data.get('pyg_time', None)
            dgl_time = data.get('dgl_time', None)
            if pyg_time is not None and dgl_time is not None and pyg_time > 0 and dgl_time > 0:
                print(f"{operator} PyG={pyg_time:.6f} DGL={dgl_time:.6f}")
    
    def _print_large_scale_table(self):
        """打印大规模性能（仅数据）"""
        if not self.large_scale_data:
            return
        print("\n大规模性能:")
        for operator in ['edge_softmax', 'gspmm', 'gsddmm']:
            data = self.large_scale_data.get(operator)
            if not data:
                continue
            total_time = data.get('total_time')
            avg_time = data.get('avg_time')
            if total_time is not None:
                if avg_time is not None:
                    print(f"{operator} total={total_time:.6f}s avg={avg_time:.6f}s")
                else:
                    print(f"{operator} total={total_time:.6f}s")
    
    def _print_test_statistics(self):
        """已弃用：不再打印统计装饰信息"""
        return
        

def run_test(test_file, test_name=None, verbose=True, data_collector=None):
    """运行单个测试文件"""
    cmd = ["python", "-m", "pytest", test_file, "-v", "-s"]
    if not verbose:
        cmd.append("-q")
    result = subprocess.run(cmd, cwd=project_root, capture_output=True, text=True)
    
    # 收集测试数据
    if data_collector and result.stdout:
        data_collector.parse_test_output(test_name or test_file, result.stdout)
    
    return result.returncode == 0

def run_all_tests():
    """运行所有算子测试"""
    # 创建数据收集器
    data_collector = TestDataCollector()
    
    # 测试文件列表
    test_files = [
        ("test_ops_edge_softmax.py", "Edge Softmax 算子测试"),
        ("test_ops_gspmm.py", "GSPMM 算子测试"), 
        ("test_ops_gsddmm.py", "GSDDMM 算子测试")
    ]
    
    for test_file, test_name in test_files:
        file_path = Path(__file__).parent / test_file
        if file_path.exists():
            success = run_test(str(file_path), test_name, data_collector=data_collector)
        else:
            print(f"❌ 测试文件不存在: {test_file}")
    # 打印精简数据
    data_collector.print_summary_tables()

def run_specific_test(operator):
    """运行特定算子的测试"""
    test_map = {
        "edge_softmax": "test_ops_edge_softmax.py",
        "gspmm": "test_ops_gspmm.py", 
        "gsddmm": "test_ops_gsddmm.py"
    }
    
    if operator not in test_map:
        print(f"❌ 未知的算子: {operator}")
        print(f"支持的算子: {', '.join(test_map.keys())}")
        return False
    
    test_file = test_map[operator]
    file_path = Path(__file__).parent / test_file
    
    if not file_path.exists():
        print(f"❌ 测试文件不存在: {test_file}")
        return False
    
    # 创建数据收集器
    data_collector = TestDataCollector()
    
    success = run_test(str(file_path), f"{operator.upper()} 算子测试", data_collector=data_collector)
    if success:
        data_collector.print_summary_tables()
    return success

def main():
    parser = argparse.ArgumentParser(description="图算子精度和性能测试")
    parser.add_argument("--operator", "-o", 
                       choices=["edge_softmax", "gspmm", "gsddmm"],
                       help="运行特定算子的测试")
    parser.add_argument("--quiet", "-q", action="store_true",
                       help="静默模式，减少输出")
    
    args = parser.parse_args()
    
    if args.operator:
        success = run_specific_test(args.operator)
        sys.exit(0 if success else 1)
    else:
        run_all_tests()

if __name__ == "__main__":
    main()
